import gym
import matplotlib.pyplot as plt
from matplotlib import animation
import os


def display_frames_as_gif(frames, id, path, frame_skip=1):
    try:
        dpi = 60
        h, w = frames[0].shape[:2]
        figsize = (w / dpi, h / dpi)
        plt.figure(figsize=figsize, dpi=dpi)
        plt.axis('off')
        patch = plt.imshow(frames[0])

        def animate(i):
            patch.set_data(frames[i * frame_skip])

        # Reduce the number of frames by skipping every `frame_skip` frames
        num_frames = len(frames) // frame_skip
        anim = animation.FuncAnimation(plt.gcf(), animate, frames=num_frames, interval=5)
        plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
        anim.save(os.path.join(path, f'{id}.gif'), writer='pillow', fps=30)
        plt.close()
    except Exception as e:
        print(f'env::display_frames_as_gif error -> {e}')

env = gym.make("HalfCheetah-v2")

render_mode="rgb_array"

# Path to save the GIFs
save_path = "gifs_mujoco_v2"
os.makedirs(save_path, exist_ok=True)

# Run three episodes
for episode in range(3):
    # Reset the environment and get the initial observation
    observation = env.reset()
    frames = []  # List to store frames for the current episode
    episode_over = False
    step_cnt = 0
    ep_r = 0

    while not episode_over:
        # Sample a random action from the action space
        action = env.action_space.sample()

        # Execute the action and get the next state, reward, and done status
        observation, reward, episode_over, info = env.step(action)
        step_cnt += 1

        # Render the current step and append the frame to the list
        # frame = env.render(mode=render_mode)
        # frames.append(frame)

        ep_r += reward

    # Save the frames as a GIF
    # display_frames_as_gif(frames, f"episode_{episode + 1}", save_path)
    print(f"Episode {episode + 1}: Total Reward = {ep_r}, Steps = {step_cnt}")

# Close the environment
env.close()